from langchain_core.documents import Document
from major.models_manager import embedding_model,chat_model

# 中文版文档
documents = [
    Document(
        page_content="狗是很好的伴侣，以忠诚和友好著称。",
        metadata={"source": "哺乳类宠物文档"},
    ),
    Document(
        page_content="猫是独立性强的宠物，通常喜欢有自己的空间。",
        metadata={"source": "哺乳类宠物文档"},
    ),
    Document(
        page_content="金鱼是适合初学者的热门宠物，饲养要求相对简单。",
        metadata={"source": "鱼类宠物文档"},
    ),
    Document(
        page_content="鹦鹉是聪明的鸟类，能够模仿人类语言。",
        metadata={"source": "鸟类宠物文档"},
    ),
    Document(
        page_content="兔子是群居动物，需要充足的空间来跳跃活动。",
        metadata={"source": "哺乳类宠物文档"},
    ),
]

from langchain_chroma import Chroma

# 1.建立文档库
vectorstore = Chroma.from_documents(
    documents,
    embedding=embedding_model.get_model(),
)

llm = chat_model.get_model()
def generate(query):
    # 2.检索
    docs = vectorstore.similarity_search(query, k=3)

    # 3.生成
    # 先把pagecontent和metadata拼接起来
    docs = ['内容:'+doc.page_content + '\n' + '来源:'+doc.metadata['source'] for doc in docs]
    #print(docs)
    doc_context = '\n\n'.join(docs)
    #print(doc_context)
    # 拼接prompt
    prompt = f'''
    根据文档回答用户问题，如果文档中不包含用户问题的答案，请回答不知道。
    文档：
    {doc_context}
    用户问题：{query}
    '''
    print(prompt)
    res = llm.stream(prompt)
    for i in res:
        print(i.content, end='')


generate('鹦鹉是什么')